from .unet import UNet
import torch as t
from torch import nn
import numpy as np
from einops import rearrange

def fill_noise(x, noise_type):
    """Fills tensor `x` with noise of type `noise_type`."""
    if noise_type == 'u':
        x.uniform_()
    elif noise_type == 'n':
        x.normal_() 
    else:
        assert False

def np_to_torch(img_np):
    '''Converts image in numpy.array to torch.Tensor.

    From C x W x H [0..1] to  C x W x H [0..1]
    '''
    return t.from_numpy(img_np)[None, :]

def get_noise(input_depth, method, spatial_size, noise_type='u', var=1./10):
    """Returns a pytorch.Tensor of size (1 x `input_depth` x `spatial_size[0]` x `spatial_size[1]`) 
    initialized in a specific way.
    Args:
        input_depth: number of channels in the tensor
        method: `noise` for fillting tensor with noise; `meshgrid` for np.meshgrid
        spatial_size: spatial size of the tensor to initialize
        noise_type: 'u' for uniform; 'n' for normal
        var: a factor, a noise will be multiplicated by. Basically it is standard deviation scaler. 
    """
    if isinstance(spatial_size, int):
        spatial_size = (spatial_size, spatial_size)
    if method == 'noise':
        shape = [1, input_depth, spatial_size[0], spatial_size[1]]
        net_input = t.zeros(shape)
        
        fill_noise(net_input, noise_type)
        net_input *= var            
    elif method == 'meshgrid': 
        assert input_depth == 2
        X, Y = np.meshgrid(np.arange(0, spatial_size[1])/float(spatial_size[1]-1), np.arange(0, spatial_size[0])/float(spatial_size[0]-1))
        meshgrid = np.concatenate([X[None,:], Y[None,:]])
        net_input=  np_to_torch(meshgrid)
    else:
        assert False
        
    return net_input

class DIP(nn.Module):
    def __init__(self,params):
        super().__init__()
        net_type = params["net_type"] # "UNET"
        img_shape = params["img_shape"] # (256,256,3)
        try:
            net_input_mode = params["net_input_mode"]
        except KeyError:
            net_input_mode = 'noise'

        if net_input_mode == 'noise':
            input_depth = 1
        elif net_input_mode == 'img':
            if len(img_shape) == 2:
                input_depth = 1
            else:
                input_depth = img_shape[2]

        self.net_input = get_noise(input_depth, "noise", img_shape[:-1]).type(dtype).to(params["device"])
        
        dtype = t.cuda.FloatTensor
        if net_type == "UNET":
            self.net = UNet(num_input_channels=input_depth, num_output_channels=img_shape[-1], 
                        feature_scale=8, more_layers=1, 
                        concat_x=False, upsample_mode='deconv', 
                        pad='zero', norm_layer=t.nn.InstanceNorm2d, need_sigmoid=True, need_bias=True)


    def forward(self,*args):
        try:
            self.net_input=args[0]
            if len(self.net_input.shape) == 2:
                self.net_input = self.net_input.unsqueeze(-1)
        except IndexError:
            pass
        return rearrange(self.net(self.net_input),'batch channel height width -> height width (channel batch)')
